import numpy as np
import cv2

def compute_exg(rgb: np.ndarray) -> np.ndarray:
    """Compute Excess Green (ExG) index from an RGB image.
    ExG = 2G - R - B  (unnormalized). Returns float32 array.
    rgb: HxWx3 in RGB order, dtype uint8 or float.
    """
    arr = rgb.astype(np.float32)
    r, g, b = arr[..., 0], arr[..., 1], arr[..., 2]
    exg = 2.0 * g - r - b
    return exg

def normalize_to_uint8(arr: np.ndarray) -> np.ndarray:
    m, M = arr.min(), arr.max()
    if M - m < 1e-6:
        return np.zeros_like(arr, dtype=np.uint8)
    out = (arr - m) / (M - m) * 255.0
    return out.clip(0, 255).astype(np.uint8)

def otsu_threshold(gray_uint8: np.ndarray) -> int:
    # Returns Otsu threshold value
    thr, _ = cv2.threshold(gray_uint8, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    return int(thr)

def exg_stress_mask(rgb: np.ndarray):
    """Return (mask, exg_uint8, thr) where mask==1 denotes 'low green' (可能压力区).
    这里用ExG的低值作为“低绿度”，因此先计算ExG->归一化->Otsu阈值，再取低于阈值的像素。"""
    exg = compute_exg(rgb)
    exg_u8 = normalize_to_uint8(exg)
    thr = otsu_threshold(exg_u8)
    # 低于阈值 => 低绿度（可能应激）
    mask = (exg_u8 < thr).astype(np.uint8)
    return mask, exg_u8, thr

def overlay_mask(rgb: np.ndarray, mask: np.ndarray, alpha: float = 0.4):
    """Overlay a red mask on top of the RGB image where mask==1."""
    overlay = rgb.copy().astype(np.float32)
    red = np.zeros_like(overlay)
    red[..., 0] = 255  # R
    blended = overlay * (1 - alpha) + red * alpha * mask[..., None]
    return blended.clip(0, 255).astype(np.uint8)